#!/Pub/Users/fuxj/.conda/envs/R_clusterProfiler/bin/Rscript
suppressPackageStartupMessages(library(argparse))
suppressPackageStartupMessages(library(tidyverse))

GO_KEGG <- function(
    geneList = NULL, saveplot = F, output_dir = "./", var_name = NULL, geneAndlogFC = NULL,
    adj_m = "BH", p_cut = .05,
    width_single = 5, height_single = 4.5, width_circ = 10, organism = "Hs") {
  if (is.null(var_name)) {
    var_name <- deparse(substitute(deg_res))
  }

  suppressPackageStartupMessages(library(cowplot))
  suppressPackageStartupMessages(library(clusterProfiler))
  suppressPackageStartupMessages(library(DOSE))
  suppressPackageStartupMessages(library(org.Hs.eg.db))
  suppressPackageStartupMessages(library(enrichplot))

  organism = tryCatch(match.arg(organism, c("Hs", "Mm")), error = \(e){
    cli::cli_alert_danger("{.var organism} must be either 'Hs' or 'Mm'.")
  })

  switch(organism,
    Hs = {
      db <- "org.Hs.eg.db"
      org_in_kegg <- "hsa"
    },
    Mm = {
      db <- "org.Mm.eg.db"
      org_in_kegg <- "mmu"
    }
  )

  # org.Hs.eg.db
  # org.Mm.eg.db

  id <- bitr(
    geneID = geneList,
    fromType = "SYMBOL",
    toType = "ENTREZID",
    OrgDb = db
  )

  KEGG_res <- enrichKEGG(
    gene = id$ENTREZID,
    organism = org_in_kegg,
    pvalueCutoff = p_cut,
    pAdjustMethod = adj_m,
    use_internal_data = F,
    qvalueCutoff = .5
  )

  if (organism == "Hs") {
    # KEGG_res <- setReadable(KEGG_res, OrgDb = get0(db), keyType = "SYMBOL")
    NULL
  }

  KEGG_bar_plot <- barplot(KEGG_res,
    title = "KEGG",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  KEGG_dot_plot <- dotplot(KEGG_res,
    title = "KEGG",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  if (nrow(as.data.frame(KEGG_res)) == 0) {
    cli::cli_alert_info(cli::style_bold("KEGG没有富集结果"))
    KEGG_bar_plot <- KEGG_dot_plot <- ggplot()
  }

  GO_MF <- enrichGO(
    gene = id$ENTREZID,
    OrgDb = db,
    keyType = "ENTREZID",
    ont = "MF", # One of "BP", "MF", and "CC" subontologies, or "ALL" for all three.
    pvalueCutoff = p_cut,
    pAdjustMethod = adj_m,
    qvalueCutoff = .5,
    readable = TRUE
  )
  # GO_MF <- setReadable(GO_MF, OrgDb = get0(db), keyType = "auto")
  # GO_MF@result %>% data.table::fwrite(., sprintf("%s_GO_MF_res.txt", var_name), sep = "\t")
  GO_MF_bar <- barplot(GO_MF,
    title = "Molecular Function",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))
  GO_MF_dot <- dotplot(GO_MF,
    title = "Molecular Function",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  if (nrow(as.data.frame(GO_MF)) == 0) {
    cli::cli_alert_info(cli::style_bold("GO MF没有富集结果"))
    GO_MF_bar <- GO_MF_dot <- ggplot()
  }

  GO_BP <- enrichGO(
    gene = id$ENTREZID,
    OrgDb = db,
    keyType = "ENTREZID",
    ont = "BP", # One of "BP", "BP", and "CC" subontologies, or "ALL" for all three.
    pvalueCutoff = p_cut,
    pAdjustMethod = adj_m,
    qvalueCutoff = .5,
    readable = TRUE
  )
  # GO_BP <- setReadable(GO_BP, OrgDb = get0(db), keyType = "ENTREZID")
  # GO_BP@result %>% data.table::fwrite(., sprintf("%s_GO_BP_res.txt", var_name), sep = "\t")
  GO_BP_bar <- barplot(GO_BP,
    title = "Biological Process",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))
  GO_BP_dot <- dotplot(GO_BP,
    title = "Biological Process",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  if (nrow(as.data.frame(GO_BP)) == 0) {
    cli::cli_alert_info(cli::style_bold("GO BP没有富集结果"))
    GO_BP_bar <- GO_BP_dot <- ggplot()
  }

  GO_CC <- enrichGO(
    gene = id$ENTREZID,
    OrgDb = db,
    keyType = "ENTREZID",
    ont = "CC", # One of "BP", "CC", and "CC" subontologies, or "ALL" for all three.
    pvalueCutoff = p_cut,
    pAdjustMethod = adj_m,
    qvalueCutoff = .5,
    readable = TRUE
  )
  # GO_CC <- setReadable(GO_CC, OrgDb = get0(db), keyType = "ENTREZID")
  # GO_CC@result %>% data.table::fwrite(., sprintf("%s_GO_CC_res.txt", var_name), sep = "\t")
  GO_CC_bar <- barplot(GO_CC,
    title = "Cellular Component",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  GO_CC_dot <- dotplot(GO_CC,
    title = "Cellular Component",
    showCategory = 10
  ) +
    theme(plot.margin = unit(c(0.2, 0.2, 0.2, 0.2), "cm"))

  if (nrow(as.data.frame(GO_CC)) == 0) {
    cli::cli_alert_info(cli::style_bold("GO CC没有富集结果"))
    GO_CC_bar <- GO_CC_dot <- ggplot()
  }

  GO_ALL <- enrichGO(
    gene = id$ENTREZID,
    OrgDb = db,
    keyType = "ENTREZID",
    ont = "ALL", # One of "BP", "CC", and "CC" subontologies, or "ALL" for all three.
    pvalueCutoff = p_cut,
    pAdjustMethod = adj_m,
    qvalueCutoff = .5,
    readable = TRUE
  )
  # GO_ALL <- setReadable(GO_ALL, OrgDb = get0(db), keyType = "ENTREZID")

  GO_ALL_bar <- barplot(GO_ALL, split = "ONTOLOGY") + facet_grid(ONTOLOGY ~ ., scale = "free")
  GO_ALL_dot <- dotplot(GO_ALL, split = "ONTOLOGY") + facet_grid(ONTOLOGY ~ ., scale = "free")

  GO_ALL_list <- list(GO_ALL_bar = GO_ALL_bar, GO_ALL_dot = GO_ALL_dot)

  GOplotIn <- KEGG_res[1:6, c(1, 2, 6, 8)]
  GOplotIn$geneID <- str_replace_all(GOplotIn$geneID, "/", ",")
  colnames(GOplotIn) <- c("ID", "Term", "adj_pval", "Genes")
  GOplotIn$Category <- "KEGG"

  if (!is.null(geneAndlogFC)) {
    genedata <- data.frame(ID = geneAndlogFC$gene, logFC = geneAndlogFC$log2FC)

    library(GOplot)
    circ <- GOplot::circle_dat(GOplotIn, genedata)
    chord <- GOplot::chord_dat(data = circ, genes = genedata)

    KEGG_circ <- GOChord(
      data = chord,
      title = "",
      space = 0.01, # Term间距
      limit = c(1, 1),
      gene.order = "logFC",
      gene.space = 0.25,
      gene.size = 3.75,
      lfc.col = c("#367EB8", "white", "#E4191C"),
      ribbon.col = brewer.pal(length(GOplotIn$Term), "Paired"),
      process.label = 10
    )

    KEGG_circ$guides$size$title <- "Terms"

    KEGG_circ$guides$size$ncol <- 1

    KEGG_circ <- KEGG_circ + theme(
      legend.position = "right",
      legend.direction = "vertical",
      legend.box = "horizontal",
      legend.box.just = "top",
      legend.justification = "center"
    )
  }

  message("Analyzing is done, begin to save plot.")

  p_list_1 <- list(
    GO_BP_dot, GO_CC_dot, GO_MF_dot,
    GO_BP_bar, GO_CC_bar, GO_MF_bar,
    KEGG_bar_plot, KEGG_dot_plot
  )

  names(p_list_1) <- c(
    "GO_BP_dot", "GO_CC_dot", "GO_MF_dot",
    "GO_BP_bar", "GO_CC_bar", "GO_MF_bar",
    "KEGG_bar_plot", "KEGG_dot_plot"
  )

  GO_dot_h <- plot_grid(
    plotlist = list(GO_BP_dot, GO_CC_dot, GO_MF_dot),
    align = "h", nrow = 1, label_size = 20 # labels = "AUTO",
  )
  GO_bar_h <- plot_grid(
    plotlist = list(GO_BP_bar, GO_CC_bar, GO_MF_bar),
    align = "h", nrow = 1, label_size = 20 # labels = "AUTO",
  )

  GO_h <- list(GO_dot_h = GO_dot_h, GO_bar_h = GO_bar_h)

  ALL_bar_p <- plot_grid(
    plotlist = list(GO_BP_bar, GO_CC_bar, GO_MF_bar, KEGG_bar_plot),
    align = "h", nrow = 2, label_size = 20 # labels = "AUTO",
  )
  ALL_dot_p <- plot_grid(
    plotlist = list(GO_BP_dot, GO_CC_dot, GO_MF_dot, KEGG_dot_plot),
    align = "h", nrow = 2, label_size = 20 # labels = "AUTO",
  )

  ALL_p <- list(ALL_bar_p = ALL_bar_p, ALL_dot_p = ALL_dot_p)


  if (saveplot) {
    if (!dir.exists(sprintf("%sGO_KEGG_%s", output_dir, var_name))) {
      dir.create(sprintf("%sGO_KEGG_%s", output_dir, var_name), showWarnings = F, recursive = T)
    } else {
      message(sprintf("Dir '%sGO_KEGG_%s' is existed.", output_dir, var_name))
    }

    KEGG_res@result %>%
      data.table::fwrite(., sprintf("%sGO_KEGG_%s/table_KEGG_%s.txt", output_dir, var_name, var_name), sep = "\t")
    GO_BP@result %>%
      data.table::fwrite(., sprintf("%sGO_KEGG_%s/table_GO_BP_%s.txt", output_dir, var_name, var_name), sep = "\t")
    GO_MF@result %>%
      data.table::fwrite(., sprintf("%sGO_KEGG_%s/table_GO_MF_%s.txt", output_dir, var_name, var_name), sep = "\t")
    GO_CC@result %>%
      data.table::fwrite(., sprintf("%sGO_KEGG_%s/table_GO_CC_%s.txt", output_dir, var_name, var_name), sep = "\t")
    GO_ALL@result %>%
      data.table::fwrite(., sprintf("%sGO_KEGG_%s/table_GO_ALL_%s.txt", output_dir, var_name, var_name), sep = "\t")


    walk(as.list(1:length(p_list_1)), function(x) {
      ggsave(
        filename = sprintf("%sGO_KEGG_%s/figure_%s.pdf", output_dir, var_name, names(p_list_1)[x]), device = "pdf",
        plot = p_list_1[[x]], width = width_single, height = height_single
      )
    })

    walk(as.list(1:length(GO_ALL_list)), function(x) {
      ggsave(
        filename = sprintf("%sGO_KEGG_%s/figure_%s.pdf", output_dir, var_name, names(GO_ALL_list)[x]), device = "pdf",
        plot = GO_ALL_list[[x]], width = width_single * 1.5, height = height_single * 3
      )
    })

    walk(as.list(1:length(GO_h)), function(x) {
      ggsave(
        filename = sprintf("%sGO_KEGG_%s/figure_%s.pdf", output_dir, var_name, names(GO_h)[x]), device = "pdf",
        plot = GO_h[[x]], width = width_single * 3, height = height_single
      )
    })

    if (length(geneAndlogFC) != 0) {
      ggsave(
        filename = sprintf("%sGO_KEGG_%s/figure_KEGG_circ_%s.pdf", output_dir, var_name, var_name), device = "pdf",
        plot = KEGG_circ, width = width_circ, height = 5
      )
    }


    walk(as.list(1:length(ALL_p)), function(x) {
      ggsave(
        filename = sprintf("%sGO_KEGG_%s/figure_%s.pdf", output_dir, var_name, names(ALL_p)[x]), device = "pdf",
        plot = ALL_p[[x]], width = width_single * 2, height = height_single * 2
      )
    })
  }

  if (length(geneAndlogFC) != 0) {
    tmp <- list(
      KEGG_res, GO_BP, GO_CC, GO_MF, GO_ALL,
      KEGG_bar_plot, KEGG_dot_plot,
      GO_BP_dot, GO_BP_bar,
      GO_CC_dot, GO_CC_bar,
      GO_MF_dot, GO_MF_bar,
      GO_ALL_dot, GO_ALL_bar,
      KEGG_circ,
      GO_dot_h, GO_bar_h,
      ALL_bar_p, ALL_bar_p
    )

    names(tmp) <- c(
      "KEGG_res", "GO_BP", "GO_CC", "GO_MF", "GO_ALL",
      "KEGG_bar_plot", "KEGG_dot_plot",
      "GO_BP_dot", "GO_BP_bar",
      "GO_CC_dot", "GO_CC_bar",
      "GO_MF_dot", "GO_MF_bar",
      "GO_ALL_dot", "GO_ALL_bar",
      "KEGG_circ",
      "GO_dot_h", "GO_bar_h",
      "ALL_bar_p", "ALL_bar_p"
    )
  } else {
    tmp <- list(
      KEGG_res, GO_BP, GO_CC, GO_MF, GO_ALL,
      KEGG_bar_plot, KEGG_dot_plot,
      GO_BP_dot, GO_BP_bar,
      GO_CC_dot, GO_CC_bar,
      GO_MF_dot, GO_MF_bar,
      GO_ALL_dot, GO_ALL_bar,
      GO_dot_h, GO_bar_h,
      ALL_bar_p, ALL_bar_p
    )

    names(tmp) <- c(
      "KEGG_res", "GO_BP", "GO_CC", "GO_MF", "GO_ALL",
      "KEGG_bar_plot", "KEGG_dot_plot",
      "GO_BP_dot", "GO_BP_bar",
      "GO_CC_dot", "GO_CC_bar",
      "GO_MF_dot", "GO_MF_bar",
      "GO_ALL_dot", "GO_ALL_bar",
      "GO_dot_h", "GO_bar_h",
      "ALL_bar_p", "ALL_bar_p"
    )
  }

  return(tmp)
}


wrp <- "直接来利用小俊姐的conda富集分析环境，对指定基因进行GO, KEGG富集分析。v2新增对小鼠的物种选项。"
parser <- ArgumentParser(
  description = wrp,
  formatter_class = "argparse.RawTextHelpFormatter"
)

parser$add_argument("-od", "--output_dir",
  action = "store",
  default = "./", metavar = "",
  help = "代表结果输出目录 ,默认当前命令执行目录，`./`"
)
parser$add_argument("-f", "--gene_file",
  action = "store",
  default = "/Pub/Users/wangyk/project/Poroject/F240108002_SKCM_sc/data/ion_channel-related_genes", metavar = "",
  help = "基因文件所在路径。要求第一列为基因symbol，利用第一列的基因symbol进行分析。"
)
parser$add_argument("-v", "--var_name",
  action = "store",
  default = "_", metavar = "",
  help = "保存文件中识别的关键字，默认为下划线"
)
parser$add_argument("-a", "--adj_p_method",
  action = "store",
  default = "BH", metavar = "",
  help = "p值矫正方法，默认BH，支持“holm”, “hochberg”, “hommel”, “bonferroni”, “BH”, “BY”, “fdr”, “none”"
)
parser$add_argument("-p", "--pvalue",
  action = "store", type = "numeric",
  default = .05, metavar = "",
  help = "结果筛选P值，默认0.05"
)
parser$add_argument("-w", "--width",
  action = "store", type = "numeric",
  default = 5, metavar = "",
  help = "单图图宽，默认5"
)
parser$add_argument("-he", "--height",
  action = "store", type = "numeric",
  default = 4, metavar = "",
  help = "单图图高，默认4"
)
parser$add_argument("-o", "--organism",
  action = "store",
  default = 'Hs', metavar = "",
  help = "默认物种类型，默认Hs，可选Mm"
)

args_in_process <- parser$parse_args()


# od check
if (!dir.exists(args_in_process$output_dir)) {
  dir.create(args_in_process$output_dir, showWarnings = FALSE, recursive = TRUE)
}

# file check
.gene_file_check <- function() {
  tryCatch(stopifnot(file.exists(args_in_process$gene_file)),
    error = function(e) {
      cli::cli_alert_danger("指定的基因文件不存在。{.path {args_in_process$gene_file}}")
    }
  )
}
.gene_file_check()

# read data
if (str_detect(tolower(basename(args_in_process$gene_file)), "\\.csv$")) {
  gs <- read.csv(args_in_process$gene_file)[[1]]
} else {
  gs <- read.delim(args_in_process$gene_file)[[1]]
}

# analyzing
a <- GO_KEGG(
  geneList = gs,
  saveplot = TRUE,
  output_dir = paste0(args_in_process$output_dir, "/"),
  var_name = args_in_process$var_name,
  geneAndlogFC = NULL,
  adj_m = args_in_process$adj_p_method,
  p_cut = args_in_process$pvalue,
  width_single = args_in_process$width,
  height_single = args_in_process$height,
  width_circ = 10,
  organism = args_in_process$organism
)